﻿#include <windows.h>
#include <stdio.h>
#include <iostream>
#include <string>

#include "..\DevMonCommon\DevMonoCommon.h"

int
Error(const std::string& ErrorString);

int
Usage();

int 
wmain(int argc, const WCHAR* argv[])
{
    if (argc < 2) {
		return Usage();
    }

	const WCHAR* &Cmd = argv[1];

	HANDLE hFile = CreateFile(L"\\\\.\\KDevMon", GENERIC_READ | GENERIC_WRITE,
		FILE_SHARE_READ, nullptr, OPEN_EXISTING, 0, nullptr);
	if (hFile == INVALID_HANDLE_VALUE) {
		return Error("Failed to open device");
	}

	DWORD dwBytes;
	if (::_wcsicmp(Cmd, L"Add") == 0) {
		if (::DeviceIoControl(hFile, IOCTL_DEVMON_ADD_DEVICE, (PVOID)argv[2],
			static_cast<DWORD>((::wcslen(argv[2]) + 1) * sizeof(WCHAR)), nullptr, 0, &dwBytes, nullptr) == FALSE) {
			return Error("Failed in add device");
		}
		printf("Add device %ws successful\n", argv[2]);
		return 0;
	}
	else if (::_wcsicmp(Cmd, L"Remove") == 0) {
		if (::DeviceIoControl(hFile, IOCTL_DEVMON_REMOVE_DEVICE, (PVOID)argv[2],
			static_cast<DWORD>((::wcslen(argv[2]) + 1) * sizeof(WCHAR)), nullptr, 0, &dwBytes, nullptr) == FALSE) {
			return Error("Failed in remove device");
		}
		printf("Remove device %ws successful\n", argv[2]);
		return 0;
	}
	else if (::_wcsicmp(Cmd, L"Clear") == 0) {
		if (::DeviceIoControl(hFile, IOCTL_DEVMON_REMOVE_ALL, nullptr, 0, nullptr, 0, nullptr, nullptr) == FALSE) {
			return Error("Failed in clear device");
		}
		printf("Remove all device successful\n");
		return 0;
	} else {
		printf("Unknown command\n");
		return Usage();
	}

    return 0;
}

int 
Error(const std::string& ErrorString)
{
	std::cout << "Error: " << ErrorString << ": " << GetLastError() << std::endl;
	::system("pause");
	return 1;
}

int
Usage()
{
	::system("pause");
	return 1;
}